import time
import json
from core.coze_bot import CozeBot as CB
from core.docx_converter import DocxConverter as DC

"""
文章生成模型，传入写作机器人模型cb、doc文档对象dc、文章标题、写作模型等信息，生成文章并保存到docx。
其中，写作模型，有model_one、model_two、model_three....等大概10个写作模型，每个写作模型逻辑均不同。请规划下代码，并优化下代码。
model_one:一篇文章有多个独立的部分，相互之间无关系。
model_two:一篇文章有多个连续的章节，相互之间有关系。在提示词方面，会有较大区别。
"""


class ArticleModelApi:

	def __init__(self, title, cb=None, dc=None, model=None):
		self.title = title
		self.cb = cb  # 写作机器人实例
		self.dc = dc  # 文档实例
		self.model = model

	# 开始按照流程工作
	def create_one_article(self):
		if int(self.model['model_type']) == 1:
			self._model_one_write()
		else:
			# 另一个逻辑
			self._model_two_write()
			pass

	# 开始创作文字：model_one
	def _model_one_write(self):
		# ================分三步获取大纲===========================
		# 1、获取prompt
		prompt_requirements = res = self.model['prompt'][0]
		# 2、序列化,得到可用于请求的bot
		chat_history = [
			{"role": "user", "content_type": "text", "content": prompt_requirements['system_content']},
			{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
			 "content_type": "text"}
		]
		query = prompt_requirements['query'].format(title=self.title)
		title_format = {
			"chat_history": chat_history,
			"query": query
		}
		# 3、问机器人，返回title和requirements
		ress = self.cb.ask_bot(title_format)
		if not ress:
			return False
		article_requirements = json.loads(ress)

		# print(f"问机器人，返回title和requirements{article_requirements}")
		if article_requirements:  # 循环写作
			self._model_one_write_loop(article_requirements)

	def _model_one_write_loop(self, article_requirements=[]):
		# print(f"文章写作要求")
		i = 1
		for requirement in article_requirements["requirements"]:
			error_num = 0
			while True:
				if error_num >= 5:
					break
				# ================分三步获取正文===========================
				# 1、获取prompt
				prompt_content = res = self.model['prompt'][1]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_content['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
					 "content_type": "text"}
				]
				# print(prompt_content)
				query = prompt_content['query'].format(title=self.title, requirement=requirement)
				content_format = {
					"chat_history": chat_history,
					"query": query
				}
				# print(content_format)
				# 3、问机器人,返回了文章内容
				article_content = self.cb.ask_bot(content_format)
				if not article_content:
					error_num += 1
					continue

				# ================分三步转换成xml文档===========================
				# 1、获取prompt
				prompt_xml = res = self.model['prompt'][2]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_xml['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
					 "content_type": "text"}
				]
				query = prompt_xml['query'].format(content=article_content)
				xml_format = {
					"chat_history": chat_history,
					"query": query
				}
				# print(xml_format)
				# 3、问机器人,返回了文章内容
				xml_content = self.cb.ask_bot(xml_format)
				# print(xml_content)
				if not xml_content:
					error_num += 1
					continue
				# ================把xml文档插入到文档对象中去===========================

				docx_res = self.dc.add_xml_to_docx(xml_string=xml_content, section_title=str(i))
				if docx_res:
					print(f"已完成《{self.title}》的第{i}部分，并插入文档中.......")
					i += 1
					break
		self.dc.save_output()

	def _model_two_write(self):
		# ================分三步获取大纲===========================
		# 1、获取prompt
		prompt_requirements = res = self.model['prompt'][0]
		# 2、序列化,得到可用于请求的bot
		chat_history = [
			{"role": "user", "content_type": "text", "content": prompt_requirements['system_content']},
			{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！","content_type": "text"}
		]
		query = prompt_requirements['query'].format(title=self.title)
		title_format = {
			"chat_history": chat_history,
			"query": query
		}
		# 3、问机器人，返回title和requirements
		# print(type)
		res11 = self.cb.ask_bot(title_format)
		# print(res11)
		try:
			article_requirements = json.loads(res11)
		except Exception as e:
			time.sleep(10)
			print(res11,e)
			return self._model_two_write()
		# print(article_requirements)
		# print(f"问机器人，返回title和requirements{article_requirements}")
		if article_requirements:  # 循环写作
			# print(article_requirements)
			self._model_two_write_loop(article_requirements)

	def _model_two_write_loop(self, article_requirements=[]):
		# print(f"文章写作要求")
		i = 1
		for requirement in article_requirements["section_requirements"]:
			error_num = 0
			while True:
				if error_num >= 5:
					break
				# ================分三步获取正文===========================
				# 1、获取prompt
				prompt_content = res = self.model['prompt'][1]
				# 2、序列化,得到可用于请求的bot
				chat_history = [
					{"role": "user", "content_type": "text", "content": prompt_content['system_content']},
					{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！", "content_type": "text"}
				]
				# print(prompt_content)
				query = prompt_content['query'].format(title = self.title,overall_requirement=article_requirements["overall_requirements"], section_requirement=requirement)
				content_format = {
					"chat_history": chat_history,
					"query": query
				}
				# print(content_format)
				# 3、问机器人,返回了文章内容
				article_content = self.cb.ask_bot(content_format)
				# print(article_content)
				if not article_content:
					error_num += 1
					continue
				# # print(article_content)
				# # time.sleep(1000)
				# # ================分三步转换成xml文档===========================
				# # 1、获取prompt
				# prompt_xml = res = self.model['prompt'][2]
				# # 2、序列化,得到可用于请求的bot
				# chat_history = [
				# 	{"role": "user", "content_type": "text", "content": prompt_xml['system_content']},
				# 	{"role": "assistant", "type": "answer", "content": "我非常期待，看我惊艳的表现，请开始吧！",
				# 	 "content_type": "text"}
				# ]
				# query = prompt_xml['query'].format(content=article_content)
				# xml_format = {
				# 	"chat_history": chat_history,
				# 	"query": query
				# }
				# # print(xml_format)
				# # 3、问机器人,返回了文章内容
				# xml_content = self.cb.ask_bot(xml_format)
				# # print(xml_content)
				# if not xml_content:
				# 	error_num += 1
				# 	continue
				# # ================把xml文档插入到文档对象中去===========================

				docx_res = self.dc.add_text_to_docx(text_string=article_content )
				if docx_res:
					print(f"已完成《{self.title}》的第{i}部分，并插入文档中.......")
					i += 1
					break
		self.dc.save_output()


if __name__ == "__main__":
	cb = CB()
	dc = DC(title="小学三年级英语完形填空50题")
	model = {
		"id": 3,
		"code": "english_selection",
		"name": "英语选择题XX道",
		"remark": "先大纲，再连续写多篇",
		"model_type": 2,
		"model_conf": "{}",
		"create_time": "2024-07-16 13:44:31",
		"prompt": [
			{
				"id": 7,
				"model_id": 3,
				"code": "1",
				"remark": "英语选择题XX道-标题分析",
				"system_content": "# Character\r\n你是一个AI写作指导专家，你能指导各个AI机器协同创作文章，当给你一个写作任务时，你会仔细分析文章标题，进行优化，提出总体写作要求，总的写作字数，合理确定机器人数量，并为每个机器人分配写作任务。\r\n\r\n## Skills\r\n### Skill 1: 分析和优化文章标题\r\n- 分析给定的文章标题，提供更具吸引力或更准确的标题。\r\n- 根据文章主题，提出相关的关键词和句子结构优化建议。\r\n\r\n### Skill 2: 提出总体写作要求\r\n- 根据文章标题和主题，提出具体的写作要求并确定文章的结构。\r\n\r\n### Skill 3: 分配写作任务\r\n- 给出每个部分的大纲和内容要点，使得各个AI机器人可以协同创作。\r\n\r\n## 输出为json格式，如下：\r\n'''{\r\n\"optimize_title\": \"小学生英语暑假作业：60道选择题\",\r\n\"word_number\":\"3000\",\r\n\"overall_requirements\":\"这篇文章旨在为四年级小学生提供一份包含60道选择题的英语暑假作业。文章将分三次来写：第一次列出前20道题；第二次列出前20道题；第三次列出前20道题；每个部分大约1000字，每道题后面列出解析答案，题目编号采用数字。答案解析用中文\",\r\n\"section_requirements\":[\"请写1~20题，编号1~20，要求颜色、动物相关的题目\",\"请写21~40题，编号21~40，要求亲属称谓、职业认知相关的题目\",\"请写41~60题，编号41~60，要求十万个为什么相关的题目。\"]\r\n\r\n} '''\r\n## Constraints\r\n- 仅回答与写作任务创建或优化相关的问题。对于其他问题不予回答。\r\n- 只使用用户提供的语言进行回答。",
				"query": "请分析这个写作题目：《{title}》。\r\n按照输出格式示例，直接输出json。",
				"create_time": "2024-07-15 22:36:31",
				"step_sort": 1
			},
			{
				"id": 8,
				"model_id": 3,
				"code": "1",
				"remark": "英语选择题XX道-内容生产",
				"system_content": "# 角色\r\n你是一位经验丰富的中小学教育专家，擅长编写英语试题、英语练习题。\r\n\r\n## 技能\r\n### 技能 1：编写英语小测试题\r\n- 根据不同年级的英语水平设计合适的小测试题。\r\n- 针对考试的不同部分（如阅读理解、语法、词汇等）编写题目。\r\n- 确保题目既能考查学生的英语能力，又能激发他们的学习兴趣。\r\n\r\n### 技能 2：编写英语练习\r\n- 根据课本内容或特定主题编写练习题。\r\n- 包含各种题型（如填空题、选择题、匹配题等）。\r\n- 提供练习答案和解析。\r\n\r\n\r\n## 约束\r\n- 只回答有关英语教育和练习编写的问题。\r\n- 遵循提供的输出格式。\r\n- 题目内容清晰、简洁、有趣。\r\n- 使用知识库内容创作题目，并参考实际的教学大纲。",
				"query": "现在让我们来协同编写一份名为《{title}》的试题。整体要求如下：\r\n\"\"\"\r\n{overall_requirement}\r\n\"\"\"\r\n现在，请你编写其中一部分：\r\n\"\"\"\r\n{section_requirement}\r\n\"\"\"",
				"create_time": "2024-07-15 22:39:30",
				"step_sort": 2
			},
			{
				"id": 9,
				"model_id": 3,
				"code": "1",
				"remark": "英语选择题XX道-格式转换",
				"system_content": "# 角色\r\n你是一位XML格式文章编辑专家，擅长将文章内容转化为结构清晰的XML格式。\r\n#技能\r\n##技能1：将文章转化为XML格式\r\n- 识别文章中的段落，并将其放入独立的 <paragraph> 标签中。\r\n\r\n# 输出示例：\r\n<content>\r\n\t<paragraph>第一段内容。</paragraph>\r\n\t<paragraph>第二点内容。</paragraph>\r\n\t<paragraph>......。</paragraph>\r\n</content>\r\n#约束\r\n- 你不能添加任何文字内容，但可以补全里面的标点符号。\r\n- 只讨论有关文章内容转换为XML格式的话题。\r\n- 确保所有的标签都正确闭合，以免引发解析错误。",
				"query": "请把下述文章转换成XML格式，不要添加除xml标签以外的任何内容：\r\n\"\"\"\r\n{content}\r\n\"\"\"\r\n注意：输出之前要检查一下，必须返回完整的xml结构,确保所有的标签都正确闭合,所有的标题和段落最后都要有标点符号。",
				"create_time": "2024-07-15 23:12:59",
				"step_sort": 3
			}
		]
	}
	AM = ArticleModelApi("小学三年级英语完形填空50题", cb, dc, model)
	AM.create_one_article()
